# coding=utf-8

import os
import sys
import logging

from PIL import Image
from io import BytesIO

import requests


def get_img(img_url, full_path_image_name):
    """
    获取一个样本图片并保存

    param img_url 将要下载的图片 url
    param fullpath_image_name 图片保存的完整路径
    """
    r = requests.get(img_url)

    if r.status_code == 200:
        img = Image.open(BytesIO(r.content))
        img.save("%s.jpg" % full_path_image_name, "JPEG")
    else:
        logging.error("获取图片失败")


def get_samples(img_url, samples_dir, sample_count):
    """
    获取指定数量的样品

    param n 样品数量
    """

    for i in range(sample_count):
        get_img(img_url, "%s/%s" % (samples_dir, i))


def binary_conversion(full_path_image_name):
    """
    二值化图片

    param full_path_image_name 指定图片进行二值化
    """
    img = Image.open(full_path_image_name)
    img = img.convert("RGBA")
    pixdata = img.load()

    for y in range(img.size[1]):
        for x in range(img.size[0]):
            if pixdata[x, y][0] < 90:
                pixdata[x, y] = (0, 0, 0, 255)
    for y in range(img.size[1]):
        for x in range(img.size[0]):
            if pixdata[x, y][1] < 136:
                pixdata[x, y] = (0, 0, 0, 255)
    for y in range(img.size[1]):
        for x in range(img.size[0]):
            if pixdata[x, y][2] > 0:
                pixdata[x, y] = (255, 255, 255, 255)

    img.save(full_path_image_name, "JPEG")


def get_fonts(full_path_image_name, fonts_dir):
    """
    切割生成 font 图片，这里需要人工干预
    """

    # crop(left, top, right, bottom)
    img = Image.open(full_path_image_name)
    for i in range(4):
        x = 16 + i*15
        y = 2
        img.crop((x, y, x+7, y+10)).save("%s_%s/%s.jpg" %
                                         (full_path_image_name.replace("/", ""), fonts_dir, i),
                                         "JPEG")


def recognize(img_url):
    """对 url 中图片进行识别"""
    pass


if __name__ == "__main__":
    # https://login.homeinns.com/home/GetVerifyCode

    samples_dir = "samples"
    fonts_dir = "fonts"

    try:
        img_url = sys.argv[1]
        sample_count = int(sys.argv[2])
    except IndexError as e:
        print("python ocr.py img_url smaple_count")
    else:
        print("fetching sample images ...")
        get_samples(img_url, samples_dir, sample_count)

        print("binary conversion images ...")
        for filename in os.listdir(samples_dir):
            if filename.endswith("jpg"):
                binary_conversion("%s/%s" % (samples_dir, filename))

        print("get fonts ...")
        for filename in os.listdir(samples_dir):
            if filename.endswith(".jpg"):
                get_fonts("%s/%s" % (samples_dir, filename), fonts_dir)
